///@file IndexSet.cpp
///@brief Omplementation of the class representing an IndexSet for polynomial chaos expansion
///@author Arnaud Duval
///@version 0.0
///@date 2010/04/14

#include "IndexSet.h"
#include "ProbUtils.h"


///@brief Default constructor
IndexSet::IndexSet()
{
	nb = 0;
}

///@brief Constructor with arguments for polynomial chaos
///@brief Build the set of indexes, of size M, with a sum less or equal p
///@param M number of input variables for polynomial chaos
///@param p maximum degree of polynomial chaos expansion
IndexSet::IndexSet(int M, int p)
{
	int tot = Binomial(M+p, p);
	nb = tot;
	set = new MultiIndex[tot];
	for(int i = 0 ; i < tot ; i++)
		set[i] = MultiIndex(M);
	int count = 0;

	///@brief Special case for first member
	for(int v = 1 ; v <= M ; v++)
		set[count][v] = 0;
	count++;

	for(int op = 1 ; op <= p ; op++)
	{	
		///@brief Special case for j=0
		for(int w = 1 ; w <= M ; w++)
		{
			for(int v = 1 ; v <= M ; v++)
			{
				if(v==w)
					(set[count])[v] = op;
				else
					(set[count])[v] = 0;
			}
			count ++;
		}
	
		for(int j = 2; (j <= M)&&(j <= op) ; j++)
		{
			MultiIndex aj(j+1);

			aj[1] = op - j + 1;
			for(int k = 2 ; k <= j ; k++)
				aj[k] = 1;
			aj[j+1] = -1;

			bool OKjuplet = false;
			while(!OKjuplet)
			{
				///@brief processing of a j-uplet
				///@brief Conversion to a M-uplet
				MultiIndex a(M);
				for(int v = 1 ; v <= M ; v++)
				{
					if(j < M)	//Si besoin de réajuster
					{
						if( v <= M-j)
							a[v] = 0;
						else
							a[v] = aj[v - M + j];
					}
					else
						a[v] = aj[v];
				}
				///@brief Sort of the obtained M-uplet
				a.BubbleSort();
				///@brief computation of all possibles permutations, L-algorithm from Knuth
				bool all = false;
				while(!all)
				{
					set[count] = a;
					count++;
					int k = M - 1;
					while((k!=0)&&(a[k] >= a[k+1]))
					{
						k--;
					}
					if(k == 0)
						all = true;
					if(!all)
					{
						int l = M;
						while(a[k] >= a[l])
							l--;
						a.Swap(l,k);
						int r = k+1;
						l = M;
						while(r < l)
						{
							a.Swap(r,l);
							r ++;
							l--;
						}
					}
				}
				if(aj[2] < aj[1] - 1)
				{
					aj[1] = aj[1] - 1;
					aj[2] = aj[2] + 1;
				}
				else
				{
				
					int k = 3;
					int s = aj[1] + aj[2] - 1;
					while((aj[k] >= aj[1] - 1)&&(k<=j))
					{
						s = s + aj[k];
						k = k + 1;
					}
					if(k > j)
					{
						OKjuplet = true;
					}
					else
					{
						int x = aj[k] + 1;
						aj[k] = x;
						k = k - 1;
						while(k > 1)
						{
							aj[k] = x;
							s = s - x;
							k = k - 1;
						}
						aj[1] = s;
					}
				}
			}
		}
	}
}

///@brief Constructor with given size
///@param size Number of multiindexes in the indexSet
IndexSet::IndexSet(unsigned int size)
{
	nb = size;
	set = new MultiIndex[nb];
}

///@brief Copy constructor
///@param s IndexSet object to copy
IndexSet::IndexSet(const IndexSet& s)
{
	nb = s.nb;
	set = new MultiIndex[nb];
	for(unsigned int i = 0 ; i < nb ; i++)
		set[i] = s.set[i];
}

///@brief Copy constructor for temporary objects
///@param s Pointer to IndexSet object to copy
IndexSet::IndexSet(IndexSet* s)
{
	nb = (*s).nb;
	set = new MultiIndex[nb];
	for(unsigned int i = 0 ; i < nb ; i++)
		set[i] = (*s).set[i];
}

///@brief Destructor
IndexSet::~IndexSet()
{
	if(nb)
		delete[] set;
}

///@brief Overload of = operator
///@param s IndexSet object to copy
IndexSet& IndexSet::operator=(const IndexSet& s)
{
	nb = s.nb;
	set = new MultiIndex[nb];
	for(unsigned int i = 0 ; i < nb ; i++)
		set[i] = s.set[i];
	return (*this);
}
///@brief Overload of = operator for temporary object
///@param s Pointer to IndexSet object to copy
IndexSet& IndexSet::operator=(IndexSet* s)
{
	nb = (*s).nb;
	set = new MultiIndex[nb];
	for(unsigned int i = 0 ; i < nb ; i++)
		set[i] = (*s).set[i];
	return (*this);
}

///@brief Cardinality of set
int IndexSet::Cardinality()
{
	return nb;
}

///@brief Overload of [] operator : direct access to a MultiIndex
///@param i Index of multiIndex to return
MultiIndex& IndexSet::operator[](unsigned int i) const
{
	assert(i >= 0);
	assert(i < nb);
	return set[i];
}

///@brief Return a subset I_i set of multi index, which i-th component is not zero
///@param i Index of the component which cannot be zero
IndexSet IndexSet::I_i(unsigned int i)
{
	unsigned int count = 0;
	for(unsigned int j = 0 ; j < nb ; j++)
	{
		if(set[j][i] != 0)
			count ++;
	}
	IndexSet res;
	res.nb = count;
	res.set = new MultiIndex[count];
	count = 0;
	for(unsigned int j = 0 ; j < nb ; j++)
	{
		if(set[j][i] != 0)
		{
			res.set[count] = set[j];
			count ++;
		}
	}
	return res;
}

///@brief Remove a MultiIndex
///@brief id Index of MultiIndex to be removed
void IndexSet::Remove(unsigned int id)
{
	assert(id >= 0);
	assert(id < nb);
	MultiIndex* temp = new MultiIndex[nb-1];
	for(unsigned int i = 0 ; i < nb-1 ; i++)
	{
		if(i < id)
			temp[i] = set[i];
		else if(i >= id)
			temp[i] = set[i+1];
	}
	nb--;
	delete[] set;
	set = new MultiIndex[nb];
	for(unsigned int i = 0 ; i < nb ; i++)
		set[i] = temp[i];
	delete[] temp;
}


///@brief Make a truncation from q-norm
///@param p truncation order
///@param q Value of q for the q norm
void IndexSet::Truncate_qNorm(unsigned int p, double q)
{
	bool OK = false;
	while(!OK)
	{
		OK = true;
		unsigned int i = 0;
		bool found = false;
		while((i < nb)&&(!found))
		{
			if(set[i].qNorm(q) > p)
			{
				found = true;
				OK = false;
				Remove(i);
			}
			i++;
		}
	}
}

///@brief Make a truncation from qw-norm
///@param p truncation order
///@param q Value of q for the q-w norm
///@param w Vector of truncation weights
void IndexSet::Truncate_qwNorm(unsigned int p, double q, Vector w)
{
	bool OK = false;
	while(!OK)
	{
		OK = true;
		unsigned int i = 0;
		bool found = false;
		while((i < nb)&&(!found))
		{
			if(set[i].qwNorm(q,w) > p)
			{
				found = true;
				OK = false;
				Remove(i);
			}
			i++;
		}
	}
}


///@brief Overload of << operator
///@param s out stream
///@param i IndexSet object to put in the out stream
std::ostream& operator<<(std::ostream &s, const IndexSet& i)
{
	for(unsigned int j = 0 ; j < i.nb ; j++)
	{
		s << i.set[j];
		if(j != i.nb - 1)
			s << "\n";
	}
	return s;
}

///@brief Overload of + operator : add an IndexSet to another one
///@param s1 left hand operand : IndexSet
///@param s2 right hand operand : IndexSet
IndexSet operator+(const IndexSet& s1, const IndexSet& s2)
{
	if(s1.nb == 0)
		return s2;
	if(s2.nb == 0)
		return s1;
	IndexSet res;
	res.nb = s1.nb;
	for(unsigned int i = 0; i < s2.nb ; i++)
	{
		bool exist = false;
		for(unsigned int j = 0 ; j < s1.nb ; j++)
		{
			if(s1.set[j] == s2.set[i])
				exist = true;
		}
		if (!exist)
			res.nb++;
	}
	res.set = new MultiIndex[res.nb];
	for(unsigned int i = 0 ; i < s1.nb ; i++)
	{
		res.set[i] = s1.set[i];
	}
	unsigned int k = s1.nb;
	for(unsigned int i = 0; i < s2.nb ; i++)
	{
		bool exist = false;
		for(unsigned int j = 0 ; j < s1.nb ; j++)
		{
			if(s1.set[j] == s2.set[i])
				exist = true;
		}
		if (!exist)
		{
			res.set[k] = s2.set[i];
			k++;
		}
	}
	return res;
}
///@brief Overload of - operator : remove an IndexSet from another one
///@param s1 left hand operand : IndexSet
///@param s2 right hand operand : IndexSet
IndexSet operator-(const IndexSet& s1, const IndexSet& s2)
{
	if(s2.nb == 0)
		return s1;
	IndexSet res;
	res.nb = s1.nb;
	for(unsigned int i = 0; i < s2.nb ; i++)
	{
		bool exist = false;
		for(unsigned int j = 0 ; j < s1.nb ; j++)
		{
			if(s1.set[j] == s2.set[i])
				exist = true;
		}
		if (exist)
			res.nb--;
	}
	res.set = new MultiIndex[res.nb];
	unsigned int k = 0;
	for(unsigned int i = 0; i < s1.nb ; i++)
	{
		bool exist = false;
		for(unsigned int j = 0 ; j < s2.nb ; j++)
		{
			if(s1.set[i] == s2.set[j])
				exist = true;
		}
		if (!exist)
		{
			res.set[k] = s1.set[i];
			k++;
		}
	}
	return res;
	
}
